/*
 * tcping.c
 *
 * Copyright (c) 2002-2008 Marc Kirchner <mail(at)marc(dash)kirchner(dot)de>
 *
 * tcping is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * tcping is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public License
 * along with ms++. If not, see <http://www.gnu.org/licenses/>.
 *
 * tcping does a nonblocking connect to test if a port is reachable.
 * Its exit codes are:
 *     -1 an error occured
 *     0  port is open
 *     1  port is closed
 *     2  user timeout
 */

#define VERSION 1.3.5

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#ifdef HAVE_STRINGS_H
#include <strings.h>
#endif
#include <unistd.h>
#include <sys/time.h>
#include <fcntl.h>
#include <arpa/inet.h>
#include <netdb.h>

#include <pthread.h>
#include <math.h>
#define NUM_THREADS	10

void *ThreadBody(void *t);
int dotcping(struct hostent *host,int port,int verbose,long timeout_sec,long timeout_usec);
void usage();

struct thread_data{
	int  id;
	struct hostent *host;
	uint8_t *ptbuf;
	int startport;
	int nport;
	int index;
	long timeout_sec;
	long timeout_usec;
};

int main (int argc, char *argv[]) {

	int verbose=1;
	int c;
	char *cptr;

	int port=0;
	struct hostent *host;
	long timeout_sec=0, timeout_usec=0;
	
	
	/*define new var*/
	int startport = 0;		// the starting port for each thread
	int port2 = 0;			// the higher value of the port range
	uint8_t range = 0;		// define if user chose ping or multiping
	int numports =0;		// the total number of ports ( range 0 ~ 50)
	int nport = 0;			// the number of port that each thread will ping 
	int odd = 0;			// the number of ports we assing to some threads
	int index = 0;			// index to access to the info_array

	pthread_t thread[NUM_THREADS];			// N Threads
	struct thread_data tdata_array[NUM_THREADS];	// N thread structures
	uint8_t *info_array;				// Array with returned info

  	pthread_attr_t attr;
 	int rc;
  	int t;
 	void *status;
	int i = 0;
	/*end define*/


	if (argc < 3)  {
		usage(argv[0]);
	}
	
	while((c = getopt(argc, argv, "qt:u:r")) != -1) {
		switch(c) {
			case 'q':
				verbose = 0;
				break;
			case 't':
				cptr = NULL;
				timeout_sec = strtol(optarg, &cptr, 10);
				if (cptr == optarg)
					usage(argv[0]);
				break;
			case 'u':
				cptr = NULL;
				timeout_usec = strtol(optarg, &cptr, 10);
				if (cptr == optarg)
					usage(argv[0]);
				break;
			/* new case: port range*/
 			case 'r':
                                range = 1;
                                break;
			default:
				usage(argv[0]);
				break;
		}
	}
	

	if ((host = gethostbyname(argv[optind])) == NULL) {
		if (verbose)
#ifdef HAVE_HSTRERROR
			fprintf(stderr, "error: %s\n", hstrerror(h_errno));
#else
			fprintf(stderr, "error: host not found");
#endif
		exit(-1);
	}

	if (argv[optind+1]) {
		cptr = NULL;
		port = strtol(argv[optind+1], &cptr, 10);
		if (cptr == argv[optind+1])
			usage(argv[0]);
	} else {
		usage(argv[0]);
	}

	
	/*** Adding range control to the 2nd port ***/	
	if(range ==1) {
                if (argv[optind+2]) {                                                  
                        cptr = NULL;
                        port2 = strtol(argv[optind+2], &cptr, 10);                      
                        if (cptr == argv[optind+2]){                                     
                                fprintf(stderr,"Port2 it's not a valid port \n");
                                usage(argv[0]);
                                }
                } else {
                        fprintf(stderr,"Port2 is missing \n");
                        usage(argv[0]);
                }
                if (port2 < port ){
                        fprintf(stderr,"Port2 must be greater than Port1 \n");
                        usage(argv[0]);
                }
	
		numports = port2-port +1;

		if(numports >50){
			fprintf(stderr,"The maximum range is 50 \n");
			usage(argv[0]);	
		}			
        }

	/* checked all user info. now starts ping */


	if(range == 0)
		return dotcping(host,port,verbose,timeout_sec,timeout_usec);
	else {
	
		/* Initialize and set thread detached attribute */
		pthread_attr_init(&attr);
	   	pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);
	
		info_array = calloc(numports,sizeof(uint8_t));
		nport = numports/ NUM_THREADS;
		odd = numports % NUM_THREADS;
		startport = port;
		

		for(t=0; t<NUM_THREADS; t++) {
			tdata_array[t].id = t;
			tdata_array[t].host = host;
			tdata_array[t].ptbuf = info_array;
			tdata_array[t].startport = startport;
			tdata_array[t].index = index;
			tdata_array[t].timeout_sec =  timeout_sec;
			tdata_array[t].timeout_usec = timeout_usec;

			index +=nport;			
			if(t < odd){
				tdata_array[t].nport = nport+1;
				startport += nport+1;
				index++;
			}
			else{
				tdata_array[t].nport = nport;
				startport += nport;
			}
		}


		for(t=0; t<NUM_THREADS; t++) {
			
			if(tdata_array[t].nport >0){
				//printf("Main: creating thread %d\n", t);
		   		rc = pthread_create(&thread[t], &attr, ThreadBody, (void *) &tdata_array[t]); 
			   	if (rc) {
					printf("ERROR; return code from pthread_create() is %d\n", rc);
					exit(-1);
				}
			}
		}

		/* Free attribute and wait for the other threads */
		pthread_attr_destroy(&attr);
		for(t=0; t<NUM_THREADS; t++) {
			if(tdata_array[t].nport >0){
				rc = pthread_join(thread[t], &status);
			      	if (rc) {
					printf("ERROR; return code from pthread_join() is %d\n", rc);
				 	exit(-1);
				}  
			}
		}

		for(i = 0;i<numports;i++){

			if(verbose)
			switch(info_array[i]){

				case 0:
					fprintf(stdout, "%s port %d open.\n", argv[optind], port+i);
					break;
				case 1:
					fprintf(stdout, "%s port %d closed.\n", argv[optind], port+i);
					break;
				case 2:
					fprintf(stdout, "%s port %d user timeout.\n", argv[optind], port+i);
					break;
				default:
					fprintf(stderr, "error: %s port %d: getsockopt: %s\n", argv[optind], port+i, strerror(errno));
					break;
			}

		}

		free(info_array);
		pthread_exit(NULL);

	}

}

void *ThreadBody(void *t)
{
	struct thread_data *my_data;
   	my_data = (struct thread_data *) t;
	int i = 0;
	for(;i<my_data->nport;i++){
		my_data->ptbuf[my_data->index + i] = dotcping(my_data->host,my_data->startport+i,
							0,my_data->timeout_sec,my_data->timeout_usec);		
	}
	pthread_exit((void*) t);
}


int dotcping(struct hostent *host,int port,int verbose,long timeout_sec,long timeout_usec){

	
	int sockfd;
	struct sockaddr_in addr;
	
	int error = 0;
	int ret;
	socklen_t errlen;
	struct timeval timeout;
	fd_set fdrset, fdwset;
	
	
	sockfd = socket (AF_INET, SOCK_STREAM, 0);
	memset(&addr, 0, sizeof(addr));

	memcpy(&addr.sin_addr, host->h_addr_list[0], host->h_length);
	addr.sin_family = host->h_addrtype; /* always AF_INET */

	addr.sin_port = htons(port);

	fcntl(sockfd, F_SETFL, O_NONBLOCK);
	if ((ret = connect(sockfd, (struct sockaddr *) &addr, sizeof(addr))) != 0) {
		if (errno != EINPROGRESS) {
#ifdef HAVE_SOLARIS
			/* solaris immediately returns ECONNREFUSED on local ports */
			if (errno == ECONNREFUSED) {
				if (verbose) 
					fprintf(stdout, "%s port %d closed.\n",inet_ntoa(addr.sin_addr) ,ntohs(addr.sin_port));
				close(sockfd);
				return(1);
			} else {
#endif	
				if (verbose)
					fprintf(stderr, "error: %s port %d: %s\n", inet_ntoa(addr.sin_addr), ntohs(addr.sin_port), strerror(errno));
				return (-1);
#ifdef HAVE_SOLARIS
			}
#endif	
		}

		FD_ZERO(&fdrset);
		FD_SET(sockfd, &fdrset);
		fdwset = fdrset;

		timeout.tv_sec=timeout_sec + timeout_usec / 1000000;
		timeout.tv_usec=timeout_usec % 1000000;

		if ((ret = select(sockfd+1, &fdrset, &fdwset, NULL, timeout.tv_sec+timeout.tv_usec > 0 ? &timeout : NULL)) == 0) {
			/* timeout */
			close(sockfd);
			if (verbose)
				fprintf(stdout, "%s port %d user timeout.\n", inet_ntoa(addr.sin_addr), ntohs(addr.sin_port));
			return(2);
		}
		if (FD_ISSET(sockfd, &fdrset) || FD_ISSET(sockfd, &fdwset)) {
			errlen = sizeof(error);
			if ((ret=getsockopt(sockfd, SOL_SOCKET, SO_ERROR, &error, &errlen)) != 0) {
				/* getsockopt error */
				if (verbose)
					fprintf(stderr, "error: %s port %d: getsockopt: %s\n", inet_ntoa(addr.sin_addr), ntohs(addr.sin_port), strerror(errno));
				close(sockfd);
				return(-1);
			}
			if (error != 0) {
				if (verbose) 
					fprintf(stdout, "%s port %d closed.\n", inet_ntoa(addr.sin_addr),ntohs(addr.sin_port));
				close(sockfd);
				return(1);
			}
		} else {
			if (verbose)
				fprintf(stderr, "error: select: sockfd not set\n");
			exit(-1);
		}
	}
	/* OK, connection established */
	close(sockfd);
	if (verbose)
		fprintf(stdout, "%s port %d open.\n", inet_ntoa(addr.sin_addr), ntohs(addr.sin_port));
	return 0;

}

void usage(char *prog) {
	fprintf(stderr, "error: Usage: %s [-q] [-t timeout_sec] [-u timeout_usec] [-r] <host> <port> <opt:port2>\n", prog);
		exit(-1);
}
	

